{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4a3b7201-75fe-4e3f-877f-ccbcd0e72a88",
   "metadata": {},
   "source": [
    "# Stock Market Data Query Engine\n",
    "\n",
    "Here we showcase our `StockMarketDataQueryEnginePack`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0548aa2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d663ba2a-c471-451a-b3be-58d136638689",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llama_pack import download_llama_pack\n",
    "\n",
    "StockMarketDataQueryEnginePack = download_llama_pack(\n",
    "    \"StockMarketDataQueryEnginePack\",\n",
    "    \"./stock_market_data_pack\",\n",
    "    # llama_hub_url=\"https://raw.githubusercontent.com/run-llama/llama-hub/jerry/fix_stock_market/llama_hub\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e94949ea-e4ab-43d2-a384-95982057cbb1",
   "metadata": {},
   "source": [
    "#### Initialize Pack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f906de73-b763-4692-9067-6ece35a9e83f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4-1106-preview\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84d407da-8a44-4de6-aa18-413fb21f8cef",
   "metadata": {},
   "outputs": [],
   "source": [
    "stock_market_data_pack = StockMarketDataQueryEnginePack(\n",
    "    [\"MSFT\", \"AAPL\", \"GOOG\", \"AMZN\", \"NVDA\", \"META\", \"TSLA\", \"CRM\", \"AMD\", \"INTC\"],\n",
    "    period=\"1mo\",\n",
    "    llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ece6f67-16ce-4232-923e-c81a0ad4aabd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<custom_module.StockMarketDataQueryEnginePack at 0x29ba892a0>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stock_market_data_pack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f7a8b03-03ee-4bac-b3ae-f036f8aa324c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>year</th>\n",
       "      <th>month</th>\n",
       "      <th>day</th>\n",
       "      <th>Open</th>\n",
       "      <th>High</th>\n",
       "      <th>Low</th>\n",
       "      <th>Close</th>\n",
       "      <th>Volume</th>\n",
       "      <th>Dividends</th>\n",
       "      <th>Stock Splits</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2023</td>\n",
       "      <td>12</td>\n",
       "      <td>19</td>\n",
       "      <td>196.160004</td>\n",
       "      <td>196.949997</td>\n",
       "      <td>195.889999</td>\n",
       "      <td>196.940002</td>\n",
       "      <td>40714100</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2023</td>\n",
       "      <td>12</td>\n",
       "      <td>20</td>\n",
       "      <td>196.899994</td>\n",
       "      <td>197.679993</td>\n",
       "      <td>194.830002</td>\n",
       "      <td>194.830002</td>\n",
       "      <td>52242800</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2023</td>\n",
       "      <td>12</td>\n",
       "      <td>21</td>\n",
       "      <td>196.100006</td>\n",
       "      <td>197.080002</td>\n",
       "      <td>193.500000</td>\n",
       "      <td>194.679993</td>\n",
       "      <td>46482500</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2023</td>\n",
       "      <td>12</td>\n",
       "      <td>22</td>\n",
       "      <td>195.179993</td>\n",
       "      <td>195.410004</td>\n",
       "      <td>192.970001</td>\n",
       "      <td>193.600006</td>\n",
       "      <td>37122800</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2023</td>\n",
       "      <td>12</td>\n",
       "      <td>26</td>\n",
       "      <td>193.610001</td>\n",
       "      <td>193.889999</td>\n",
       "      <td>192.830002</td>\n",
       "      <td>193.050003</td>\n",
       "      <td>28919300</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>2023</td>\n",
       "      <td>12</td>\n",
       "      <td>27</td>\n",
       "      <td>192.490005</td>\n",
       "      <td>193.500000</td>\n",
       "      <td>191.089996</td>\n",
       "      <td>193.149994</td>\n",
       "      <td>48087700</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2023</td>\n",
       "      <td>12</td>\n",
       "      <td>28</td>\n",
       "      <td>194.139999</td>\n",
       "      <td>194.660004</td>\n",
       "      <td>193.169998</td>\n",
       "      <td>193.580002</td>\n",
       "      <td>34049900</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>2023</td>\n",
       "      <td>12</td>\n",
       "      <td>29</td>\n",
       "      <td>193.899994</td>\n",
       "      <td>194.399994</td>\n",
       "      <td>191.729996</td>\n",
       "      <td>192.529999</td>\n",
       "      <td>42628800</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>187.149994</td>\n",
       "      <td>188.440002</td>\n",
       "      <td>183.889999</td>\n",
       "      <td>185.639999</td>\n",
       "      <td>82488700</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>184.220001</td>\n",
       "      <td>185.880005</td>\n",
       "      <td>183.429993</td>\n",
       "      <td>184.250000</td>\n",
       "      <td>58414500</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>182.149994</td>\n",
       "      <td>183.089996</td>\n",
       "      <td>180.880005</td>\n",
       "      <td>181.910004</td>\n",
       "      <td>71983600</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>5</td>\n",
       "      <td>181.990005</td>\n",
       "      <td>182.759995</td>\n",
       "      <td>180.169998</td>\n",
       "      <td>181.179993</td>\n",
       "      <td>62303300</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>8</td>\n",
       "      <td>182.089996</td>\n",
       "      <td>185.600006</td>\n",
       "      <td>181.500000</td>\n",
       "      <td>185.559998</td>\n",
       "      <td>59144500</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>9</td>\n",
       "      <td>183.919998</td>\n",
       "      <td>185.149994</td>\n",
       "      <td>182.729996</td>\n",
       "      <td>185.139999</td>\n",
       "      <td>42841800</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>10</td>\n",
       "      <td>184.350006</td>\n",
       "      <td>186.399994</td>\n",
       "      <td>183.919998</td>\n",
       "      <td>186.190002</td>\n",
       "      <td>46792900</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>11</td>\n",
       "      <td>186.539993</td>\n",
       "      <td>187.050003</td>\n",
       "      <td>183.619995</td>\n",
       "      <td>185.589996</td>\n",
       "      <td>49128400</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "      <td>186.059998</td>\n",
       "      <td>186.740005</td>\n",
       "      <td>185.190002</td>\n",
       "      <td>185.919998</td>\n",
       "      <td>40444700</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>16</td>\n",
       "      <td>182.160004</td>\n",
       "      <td>184.259995</td>\n",
       "      <td>180.929993</td>\n",
       "      <td>183.630005</td>\n",
       "      <td>65603000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>17</td>\n",
       "      <td>181.270004</td>\n",
       "      <td>182.929993</td>\n",
       "      <td>180.300003</td>\n",
       "      <td>182.679993</td>\n",
       "      <td>47317400</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>2024</td>\n",
       "      <td>1</td>\n",
       "      <td>18</td>\n",
       "      <td>186.089996</td>\n",
       "      <td>189.139999</td>\n",
       "      <td>185.830002</td>\n",
       "      <td>188.630005</td>\n",
       "      <td>77921000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    year  month  day        Open        High         Low       Close  \\\n",
       "0   2023     12   19  196.160004  196.949997  195.889999  196.940002   \n",
       "1   2023     12   20  196.899994  197.679993  194.830002  194.830002   \n",
       "2   2023     12   21  196.100006  197.080002  193.500000  194.679993   \n",
       "3   2023     12   22  195.179993  195.410004  192.970001  193.600006   \n",
       "4   2023     12   26  193.610001  193.889999  192.830002  193.050003   \n",
       "5   2023     12   27  192.490005  193.500000  191.089996  193.149994   \n",
       "6   2023     12   28  194.139999  194.660004  193.169998  193.580002   \n",
       "7   2023     12   29  193.899994  194.399994  191.729996  192.529999   \n",
       "8   2024      1    2  187.149994  188.440002  183.889999  185.639999   \n",
       "9   2024      1    3  184.220001  185.880005  183.429993  184.250000   \n",
       "10  2024      1    4  182.149994  183.089996  180.880005  181.910004   \n",
       "11  2024      1    5  181.990005  182.759995  180.169998  181.179993   \n",
       "12  2024      1    8  182.089996  185.600006  181.500000  185.559998   \n",
       "13  2024      1    9  183.919998  185.149994  182.729996  185.139999   \n",
       "14  2024      1   10  184.350006  186.399994  183.919998  186.190002   \n",
       "15  2024      1   11  186.539993  187.050003  183.619995  185.589996   \n",
       "16  2024      1   12  186.059998  186.740005  185.190002  185.919998   \n",
       "17  2024      1   16  182.160004  184.259995  180.929993  183.630005   \n",
       "18  2024      1   17  181.270004  182.929993  180.300003  182.679993   \n",
       "19  2024      1   18  186.089996  189.139999  185.830002  188.630005   \n",
       "\n",
       "      Volume  Dividends  Stock Splits  \n",
       "0   40714100        0.0           0.0  \n",
       "1   52242800        0.0           0.0  \n",
       "2   46482500        0.0           0.0  \n",
       "3   37122800        0.0           0.0  \n",
       "4   28919300        0.0           0.0  \n",
       "5   48087700        0.0           0.0  \n",
       "6   34049900        0.0           0.0  \n",
       "7   42628800        0.0           0.0  \n",
       "8   82488700        0.0           0.0  \n",
       "9   58414500        0.0           0.0  \n",
       "10  71983600        0.0           0.0  \n",
       "11  62303300        0.0           0.0  \n",
       "12  59144500        0.0           0.0  \n",
       "13  42841800        0.0           0.0  \n",
       "14  46792900        0.0           0.0  \n",
       "15  49128400        0.0           0.0  \n",
       "16  40444700        0.0           0.0  \n",
       "17  65603000        0.0           0.0  \n",
       "18  47317400        0.0           0.0  \n",
       "19  77921000        0.0           0.0  "
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "modules = stock_market_data_pack.get_modules()\n",
    "# AAPL\n",
    "modules[\"stocks market data\"][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d1c0d2b-bbfc-449d-881e-5881913beea3",
   "metadata": {},
   "source": [
    "## Try Out Some Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8316b9b0-73d4-4749-901a-f3cb621ebfc7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;34mRetrieving with query id None: What is the average closing price for MSFT?\n",
      "\u001b[0m\u001b[1;3;38;5;200mRetrieved node with id, entering: pandas0\n",
      "\u001b[0m\u001b[1;3;34mRetrieving with query id pandas0: What is the average closing price for MSFT?\n",
      "\u001b[0m\u001b[1;3;32mGot response: 377.45850067138673\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = stock_market_data_pack.run(\"What is the average closing price for MSFT?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6db0bce5-7f0a-47b2-857e-8dc06610dd49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;34mRetrieving with query id None: What is AAPL's trading volume on the day after Christmas?\n",
      "\u001b[0m\u001b[1;3;38;5;200mRetrieved node with id, entering: pandas1\n",
      "\u001b[0m\u001b[1;3;34mRetrieving with query id pandas1: What is AAPL's trading volume on the day after Christmas?\n",
      "\u001b[0m\u001b[1;3;32mGot response: 28919300\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = stock_market_data_pack.run(\n",
    "    \"What is AAPL's trading volume on the day after Christmas?\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_hub",
   "language": "python",
   "name": "llama_hub"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
